import requests
import json
from web_spider2 import *
import json
import re

# Ollama API 地址
OLLAMA_API_URL = "http://localhost:11434"
# MODEL_NAME = "deepseek-r1:1.5b"
MODEL_NAME = "deepseek-r1-lvkong-v3"

# 系统提示
SYSTEM = """
你是一个严谨的AI助手，必须用中文回答问题。
你的回答必须详细，准确，并且遵循以下格式：
User: [用户输入]
Assistant: [你的回答]
如果用户的问题不明确，请要求用户提供更多信息。
如果问题需要实时信息，请使用联网搜索功能获取最新数据。
"""

# 模板
def build_prompt(system_prompt, user_input, search_results=None):
    prompt = f"""<|system|>
{system_prompt}<|end|>
<|user|>
User: {user_input}<|end|>
<|assistant|>
Assistant:"""
    if search_results:
        prompt += f"\n以下是根据你的问题搜索到的实时信息：\n{search_results}\n"
    return prompt

# 联网搜索函数
def search_web(query):
    return None

# 用于提取汉字和标点符号
def extract_chinese_from_string(input_string):
    """
    提取字符串中的汉字和标点符号。
    
    :param input_string: 原始字符串
    :return: 提取的汉字和标点符号字符串
    """
    # 使用正则表达式提取所有汉字和中文标点符号
    chinese_chars_and_punctuation = re.findall(r'[\u4e00-\u9fa5\u3000-\u303f\uff00-\uffef]', input_string)
    # 将列表合并为字符串
    chinese_text = ''.join(chinese_chars_and_punctuation)
    # 返回汉字和标点符号
    return chinese_text

# 主函数
def get_deepseek_response(user_input):

    if user_input == "退出" or user_input == "结束":
        continue_flag = False  # 退出循环
        print("感谢使用！")
        return

    # 判断是否需要联网搜索
    need_search = "最新" in user_input or "实时" in user_input or "今天" in user_input  # 根据关键词判断
    search_results = None
    if need_search:
        print("正在联网搜索...")
        search_results = search_web(user_input)

    # 组织请求数据
    data = {
        "model": MODEL_NAME,
        "prompt": build_prompt(SYSTEM, user_input, search_results),  # 使用模板构建完整的提示
        "stream": True  # 启用流式输出
    }

    try:
        # 发送 POST 请求
        response = requests.post(f"{OLLAMA_API_URL}/api/generate", json=data, timeout=30, stream=True)

        # 处理流式响应
        if response.status_code == 200:
            full_response = ""
            skip_flag = False  # 跳过 <think> 标签
            for line in response.iter_lines():
                if line:
                    # 解析 JSON
                    chunk = line.decode("utf-8")  # 将字节流转换为字符串
                    json_data = json.loads(chunk)  # 将字符串转换为字典

                    # print(json_data)

                    if json_data["response"] == "<think>" or json_data["response"] == "</think>":
                        skip_flag = not skip_flag
                        continue

                    if "response" in json_data and not skip_flag:
                        full_response+=json_data["response"]

                    if json_data.get("done", False):
                        break  # 如果 done 为 True，结束循环
            
            print(extract_chinese_from_string(full_response))
            print('\n')

        else:
            print(f"请求失败，状态码：{response.status_code}")
            print("API 返回：", response.text)


    except requests.exceptions.Timeout:
        print("请求超时，请检查网络或 API 服务是否正常。")
    except requests.exceptions.RequestException as e:
        print(f"请求发生错误：{e}")
        

if __name__ == '__main__':

    while True:
        input_text = input("请输入问题：")
        if input_text=="":
            break
        get_deepseek_response(input_text)